import psycopg2
from django.db.utils import load_backend, DatabaseError
from .introspection import TenantSchemaDatabaseIntrospection
# https://gitlab.com/-/snippets/2137539


psql_backend = load_backend('django.db.backends.postgresql')


class DatabaseWrapper(psql_backend.DatabaseWrapper):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        self.schemas  = ('public',)
        self.introspection = TenantSchemaDatabaseIntrospection(self)

        self.is_search_path_set = False
    
    @property
    def primary_schema(self):
        return self.schemas[0]
    
    def close(self):
        # 无聊函数1
        super().close()
        self.is_search_path_set = False

    def rollback(self):
        # 无聊函数2
        super().rollback()
        self.is_search_path_set = False

    def set_schema(self, *schemas, include_public=True):
        self.schemas = schemas

        if include_public and 'public' not in schemas:
            self.schemas += ('public',)
        
        self.is_search_path_set = False

    def _cursor(self, name=None):
        """
        所有Django发出的数据库操作指令，都需要先调用此函数以获得数据库的操作对象

        因而我们提前地加上sert search_path 命令，保证所有的操作都指向当前锁定的schema
        """
        cursor = super()._cursor(name=name)

        if self.is_search_path_set:
            return cursor
        
        try:
            target_path = ','.join(self.schemas)
            search_path_str = f'SET search_path = {target_path}'
            cursor.execute(search_path_str)
 
        except (DatabaseError, psycopg2.InternalError):
            self.is_search_path_set = False
        else:
            self.is_search_path_set = True
        
        return cursor


